import json
import torch
import os
from vllm import LLM, SamplingParams
import vllm
from tqdm import tqdm
print("vllm version:", vllm.__version__)
print("vllm location:", vllm.__file__)
import re

import time 
if __name__ == '__main__':
    sampling_params = SamplingParams(temperature=0.1, max_tokens=800)
    
    available_gpus = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
    print(available_gpus)
    llm = LLM(
            #model="/mnt/diskhd/Backup/DownloadModel/Qwen2.5-72B-Instruct-GPTQ-Int4/", 
            model="/mnt/disk2/home/wujianfeng/LLaMA-Factory/testtools/",
            tensor_parallel_size=len(available_gpus), 
            dtype=torch.bfloat16, 
            gpu_memory_utilization=0.95, 
            #kv_cache_dtype="fp8",
            max_model_len=6000, 
            enable_prefix_caching=True
            )
    
    answers = []
    data = []
    with open(f"test.jsonl", "r", encoding='utf-8') as f:
        for line in f:
            data.append(json.loads(line))
            

    print(len(data))
    for item in tqdm(data[:]):
        tools = item["apis"]

        item["targets"] = []

        #message = [{"role": "system", "content": "回答时结合上下文判断：1.当前问题是否需要使用下面的工具？\n2.上下文中是否给全了工具调用所需的所有参数？\n给出具体的思考过程。\n\n如果条件1和2都已满足请给出正确的工具调用，否则回答不需要使用工具。"}]
        message = []
        for question in item["user_messages"]:
            #question = '"' + question + '"\n\n上面这个问题是否需要system中的某个工具？上下文会话中是否给全了工具调用所需的所有参数？给出理由。如果需要请再给出工具调用，否则只回答不需要使用工具。'
            #question = "结合上下文回答：" + question
            question = '"' + question + '"\n\n判断：条件1.上面的问题是否需要使用system中的某个工具？\n条件2.上下文中是否给全了工具调用所需的所有参数？\n给出思考过程。\n\n如果条件1和条件2都已满足请给出正确的工具调用，否则回答不需要使用工具。'

            message.append({"role": "user",  "content": question})
            outputs = llm.chat(message, sampling_params=sampling_params, tools=tools)
            print("outputs: ", outputs, "\n\n\n")
            answer = outputs[0].outputs[0].text
            item["targets"].append(answer)
            #print("answer: ", answer)
            message.append({"role": "assistant",  "content": answer})

        item["message"] = message
        answers.append(item)
    
        if len(answers) % 50 == 10:
            with open("answer_7B.jsonl", "w", encoding='utf-8') as f:
                for item in answers:
                    f.write(json.dumps(item, ensure_ascii=False) + "\n")

    with open("answer_7B.jsonl", "w", encoding='utf-8') as f:
        for item in answers:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")
    
    


 